import numpy as np
import sys
import cv2
import easyocr
import math
from PIL import ImageFont, ImageDraw, Image
from matplotlib import pyplot as plt
import pprint



class TextExtractor():
    ocr_reader = easyocr.Reader(['ch_sim', 'en'])
    horizontal_list, free_list = None, None

    # Find the min & max of X & Y
    # Input: A list coordinates of the corners of an area, notice that they may NOT of a rectangle.
    #    [ [x_0, y_0]], [x_1, y_1], [x_2, y_2], [x_3, y_3] ]
    # Output: [x_left, y_lower, x_right, y_upper]
    def find_bounding_box(self, area: list) -> list:
        x_left, y_left = sys.maxsize, sys.maxsize
        x_upper, y_upper = 0, 0

        # result_item[0] is a list of 4 coordinates.
        for corner in area:
            if corner[0] <= x_left:
                x_left = int(corner[0])
            elif corner[0] >= x_upper:
                x_upper = int(corner[0])

            if corner[1] <= y_left:
                y_left = int(corner[1])
            elif corner[1] >= y_upper:
                y_upper = int(corner[1])

        return [x_left, y_left, x_upper, y_upper]


    def draw_bounding_box(self, img, weight, x_left, y_lower, x_right, y_upper):
        cv2.rectangle(img, (x_left, y_lower), (x_right, y_upper),
                      (250, 0, 0), weight)
        return img

    def draw_contour(self, img, weight, pts):
        # Polygon corner points coordinates
        pts = np.array(pts, np.int32)
        pts = pts.reshape((-1, 1, 2))

        # the last line connected to the first line.
        isClosed = True

        # Green color in BGR
        color = (0, 250, 0)

        # Using cv2.polylines() method
        # Draw a Blue polygon with
        # thickness of 1 px
        img = cv2.polylines(img, [pts], isClosed, color, weight)
        return img


    # Input: area is a list of the coordinates of the four corners.
    # Output: cropped image.
    def crop_image(self, area, img):

        # bounding_box: [x_left, y_lower, x_right, y_upper]
        bounding_box = self.find_bounding_box(area)
        # crop_img = img[y:y+h, x:x+w]
        cropped_img = img[bounding_box[1]:bounding_box[1]+(bounding_box[3] - bounding_box[1]),
                      bounding_box[0]:bounding_box[0]+(bounding_box[2] - bounding_box[0])]

        # cv2.imshow('cropped_img', cropped_img)
        # cv2.waitKey(0)
        return cropped_img



    def rotate_image(self, image, angle):
        image_center = tuple(np.array(image.shape[1::-1]) / 2)
        rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
        result = cv2.warpAffine(image, rot_mat, image.shape[1::-1], flags=cv2.INTER_LINEAR)
        return result


    def draw_Chinese_text(self, image, text, weight, x_left, y_lower):
        fontpath = "fonts/simsun.ttc"  # 宋体字集
        font = ImageFont.truetype(fontpath, weight)
        img_pil = Image.fromarray(image)
        draw = ImageDraw.Draw(img_pil)
        r, g, b, a = 250, 0, 0, 0
        draw.text((x_left, y_lower - weight), text, font=font, fill=(r, g, b, a))
        image = np.array(img_pil)
        return image

    def display_image(self, image, title):
        plt.imshow(image)
        plt.title(title)
        plt.show(block=False)

        keyboardClick = False
        while keyboardClick != True:
            keyboardClick = plt.waitforbuttonpress()
        plt.close()

        """   
        # 无标尺，不可以放大缩小，功能单薄，所以不用。
        cv2.imshow(title, image)
        cv2.waitKey(0)
        """

    # Input: an image,
    # Output: a list of bounding boxes.
    #    [ [x_left, y_lower, x_right, y_upper], ... ]
    def detect_text_boxes(self, img) -> list:
        if self.horizontal_list is None or self.free_list is None:
            (self.horizontal_list, self.free_list) = self.ocr_reader.detect(img)
        # print(" >> The horizontal_list from detect: ")
        # pprint.pprint(self.horizontal_list)
        # print("\n >> The free_list from detect: ")
        # pprint.pprint(self.free_list)

        box_list = []
        for horizon in self.horizontal_list:
            for box in horizon:
                box_int = [int(box[0]), int(box[2]), int(box[1]), int(box[3])]
                box_list.append(box_int)

        for region in self.free_list:
            for area in region:
                """
                # draw_contour(self, img, weight, pts):
                self.draw_contour(img, 5, area)
                """

                box_int = self.find_bounding_box(area)
                box_list.append(box_int)

        return box_list



    # Input: 不规则四边形的四个顶点的坐标。
    # Output: [max_width, max_height],
    # 不规则四边形中，最宽的边，和最高的边，以及开始的顶点序号。特例是正方形，正菱形。
    def calc_area_size(self, area) -> list:
        max_width, max_height = 0.0, 0.0
        max_width_start, max_height_start = -1, -1
        widths, heights = [0] * 4, [0] * 4

        for i in range(4):
            if i + 1 >= 4:
                j = 0
            else:
                j = i + 1

            widths[i] = abs(float(area[i][0] - area[j][0]))
            heights[i] = abs(float(area[i][1] - area[j][1]))

            if widths[i] > max_width:
                max_width = widths[i]
                max_width_start = i
            if heights[i] > max_height:
                max_height = heights[i]
                max_height_start = i

        max_width = math.sqrt(float(widths[max_width_start]) ** 2 + float(heights[max_width_start]) ** 2)
        max_height = math.sqrt(float(widths[max_height_start]) ** 2 + float(heights[max_height_start]) ** 2)
        return [max_width, max_height]


    # Input: area, 不规则四边形。
    # Output: 给输入的 area 的每一个顶点，打个标签。
    #   标签有四个，left_lower, left_upper, right_upper, right_lower.
    def calc_area_locations(self, area) -> list:
        x_s = [corner[0] for corner in area]
        y_s = [corner[1] for corner in area]
        x_s.sort()
        y_s.sort()

        labels = [""] * 4
        for i in range(4):
            corner = area[i]
            x_sort = x_s.index(corner[0])
            y_sort = y_s.index(corner[1])
            if x_sort < 2:
                labels[i] = "left_"
            else:
                labels[i] = "right_"

            if y_sort < 2:
                labels[i] += "lower"
            else:
                labels[i] += "upper"

        return labels


    # Input: area 是一个不规则四边形的四个顶点的坐标。
    # Output: 经过 Affine Transformation 以后，area 的四个顶点的新坐标。
    def calc_affine_destination(self, area) -> list:
        # Step 1. 先确定 Affine Transformation 完成后，新的 bounding box 的宽和高。
        [max_width, max_height] = self.calc_area_size(area)
        # print(" >> area size: [max_width, max_height] ")
        # pprint.pprint([max_width, max_height])

        # Step 2. 给输入的 area 的每一个顶点，打个标签。
        # 标签有四个，left_lower, left_upper, right_upper, right_lower.
        labels = self.calc_area_locations(area)
        # print(" >> area: ")
        # pprint.pprint(area)
        # print(" >> labels: ")
        # pprint.pprint(labels)

        # Step 3. 根据输入的 area 的顶点的顺序，输出新的 bounding box 的四个顶点。
        affine_dest = [[0.0] * 2] * 4
        [x_left, y_left, x_upper, y_upper] = self.find_bounding_box(area)
        # print(" >> bounding box: ")
        # pprint.pprint([x_left, y_left, x_upper, y_upper])

        for i in range(4):
            # print(f"   labels[{i}] == {labels[i]}")
            if labels[i] == "left_lower":
                corner = [0.0] * 2
                corner[0] = float(x_left)
                corner[1] = float(y_left)
                affine_dest[i] = corner
                continue
            elif labels[i] == "left_upper":
                corner = [0.0] * 2
                corner[0] = float(x_left)
                corner[1] = float(y_left + max_height)
                affine_dest[i] = corner
                continue
            elif labels[i] == "right_upper":
                corner = [0.0] * 2
                corner[0] = float(x_left + max_width)
                corner[1] = float(y_left + max_height)
                affine_dest[i] = corner
                continue
            elif labels[i] == "right_lower":
                corner = [0.0] * 2
                corner[0] = float(x_left + max_width)
                corner[1] = float(y_left)
                affine_dest[i] = corner
                continue
            else:
                corner = area[i]
                print(f" >> Error: {corner} is mistakenly labelled as {labels[i]}")
                break

        # print(" >> affine_dest: ")
        # pprint.pprint(affine_dest)
        return affine_dest


    def affine_transform(self, area, cropped_img):
        # Step 1. Find the 4 corner's destination coordinates  affine transformation.
        affine_dest = self.calc_affine_destination(area)
        [x_left, y_left, x_upper, y_upper] = self.find_bounding_box(area)
        # print(" >> affine_dest: ")
        # pprint.pprint(affine_dest)

        # Step 2.
        pts1 = np.float32([area[0], area[1], area[2]])
        pts2 = np.float32([affine_dest[0], affine_dest[1], affine_dest[2]])

        for i in range(3):
            pts1[i][0] -= x_left
            pts1[i][1] -= y_left
            pts2[i][0] -= x_left
            pts2[i][1] -= y_left

        # print(" >> pts1: ")
        # pprint.pprint(pts1)
        # print(" >> pts2: ")
        # pprint.pprint(pts2)

        # Step 3.
        rows, cols, ch = cropped_img.shape
        M = cv2.getAffineTransform(pts1, pts2)

        try:
            affined_img = cv2.warpAffine(cropped_img, M, (cols, rows))
        except:
            affined_img = cropped_img

        # self.display_image(cropped_img, "cropped_img")
        # self.display_image(affined_img, "affined_img")

        return affined_img



    # Input: Usually the input image is an adjusted image rectangular crop.
    # Output: A list of text: str.
    def recognize_text(self, img) -> dict:
        text_list = []

        if self.horizontal_list is None or self.free_list is None:
            (self.horizontal_list, self.free_list) = self.ocr_reader.detect(img)

        # horizontal_list is a list of [x_min, x_max, y_min, y_max]
        for horizon in self.horizontal_list:
            for box in horizon:
                # crop_img = img[y:y+h, x:x+w]
                h = int(box[3] - box[2])
                w = int(box[1] - box[0])
                crop_img = img[int(box[2]):int(box[2]) + h, int(box[0]):int(box[0] + w)]

                try:
                    results = self.ocr_reader.readtext(crop_img, detail=1, paragraph=False)

                    text_segs = []
                    for res in results:
                        if res[2] > 0.4:
                            text_segs.append(res[1])
                    text = " ".join(text_segs)
                    text_list.append(text.strip())

                    # print(" >> The results of horizontal reading text: ")
                    # pprint.pprint(results)
                except:
                    text_list.append("")


        # 1. crop the image fragment,
        # 2. detect the tilted angle,
        # 3. transform the image fragment with the angle,
        # 4. recognize the text from the rotated image fragment.
        for region in self.free_list:
            for area in region:
                cropped_img = self.crop_image(area, img)
                affined_img = self.affine_transform(area, cropped_img)

                try:
                    results = self.ocr_reader.readtext(affined_img, detail=1, paragraph=False)
                    # results = self.ocr_reader.readtext(cropped_img, detail=1, paragraph=False)
                    # print(" >> The results of free-list reading text: ")
                    # pprint.pprint(results)

                    text_segs = []
                    for res in results:
                        if res[2] > 0.0:
                            text_segs.append(res[1])
                    text = " ".join(text_segs)
                    text_list.append(text.strip())

                except:
                    text_list.append("")

        return text_list


def testrun_text_extractor():
    SOURCE, RESULT = "images/", "results/"
    img_file1 = "佛山乒乓比赛.jpeg"
    img_file2 = "钱大妈.jpeg"
    img_file3 = "核酸业务通知.JPG"


    img_filename = img_file3
    image = cv2.imread(SOURCE + img_filename)
    if image is None:
        print(f" !! Cannot read the image {img_filename} ")
        exit()

    text_extractor = TextExtractor()
    box_list = text_extractor.detect_text_boxes(image)
    text_list = text_extractor.recognize_text(image)
    print(" >> Text extraction result: ")
    pprint.pprint(text_list)

    for i in range(len(box_list)):
        # def draw_bounding_box(self, image, x_left, y_lower, x_right, y_upper):
        if i < len(box_list) - len(text_extractor.free_list):
            image = text_extractor.draw_bounding_box(image, 5, box_list[i][0], box_list[i][1], box_list[i][2], box_list[i][3])

        # draw_Chinese_text(self, image, text, x_left, y_lower):
        if i < len(text_list):
            image = text_extractor.draw_Chinese_text(image, text_list[i], 50, box_list[i][0], box_list[i][1])

    for region in text_extractor.free_list:
        for area in region:
            # draw_contour(self, img, weight, pts):
            text_extractor.draw_contour(image, 5, area)


    text_extractor.display_image(image, "Image with text recognition")
    cv2.imwrite(RESULT + img_filename, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))


if __name__ == "__main__":
    testrun_text_extractor()

